package com.cx.common.utils;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.lang.Assert;

import java.util.*;
import java.util.function.BiConsumer;
import java.util.function.Function;
import java.util.stream.Collectors;
import java.util.stream.Stream;

/**
 * @param <I> input 输入的数据类型
 * @param <O> output输出的数据类型
 * @author cx
 * @description: 树转化工具
 * @code 示例：菜单转成菜单树（类省略无参构造函数、全参数构造函数以及set、get方法）
 * <blockquote><pre>
 * class Menu{
 *     private Long id;               // 主键
 *     private String name;           // 菜单名称
 *     private Long parentId;         // 父id
 *     private Long sort;             // 排序字段
 * }
 * class MenuTree{
 *     private Long id;               // 树节点id
 *     private String name;           // 树节点名称
 *     private List<MenuTree> child;  // 子节点
 *     public MenuTree(Menu menu) {
 *         this.id = menu.getId();
 *         this.name = menu.getName();
 *     }
 * }
 * List< Menu> list=new ArrayList<>();
 * list.add(new Menu(1L, "菜单1", 0L,2L));
 * list.add(new Menu(2L, "菜单2", 1L,1L));
 * List< MenuTree> menuTreeList = new TreeGenerateUtil< Menu, MenuTree>(list)
 *     .keyMapFunction(Menu::getId)
 *     .parentKeyMapFunction(Menu::getParentId)
 *     .objectTransformFunction(MenuTree::new)
 *     .setChildFunction(MenuTree::setChild)
 *     .sortFunction(Comparator.comparingLong(Menu::getSort))
 *     .generate(0L);
 * </pre></blockquote>
 * @date 2023/8/2 14:34
 */
public class TreeGenerateUtil<I, O> {

    /**
     * 转化为树的原始list
     */
    private final List<I> list;

    /**
     * 数据对象I转化为树对象O的函数
     */
    private Function<? super I, ? extends O> objectTransform;

    /**
     * 获取数据对象I中的主键函数 一般为id
     */
    private Function<? super I, Object> keyMap;

    /**
     * 获取数据对象I中的标识层级关系的键函数 parentKey指向key 表示该节点属于哪个节点
     */
    private Function<? super I, Object> parentKeyMap;

    /**
     * 设置O对象中的子层级函数 子层级赋值给O对象的那个字段
     */
    private BiConsumer<O, List<O>> setChild;

    /**
     * I对象中排序比较函数
     */
    private Comparator<? super I> sort;

    /**
     * 保存处理过的key 避免循环层级导致的栈溢出
     */
    private final Set<Object> handleKey;

    /**
     * 构造函数 初始化原始数据集合
     *
     * @param list 原始数据集合
     */
    public TreeGenerateUtil(List<I> list) {
        this.list = list;
        this.handleKey = new HashSet<>();
    }

    /**
     * 设置 数据对象I转化为树对象O的函数
     *
     * @param objectTransform 数据对象I转化为树对象O的函数
     * @return 返回this用于链式调用
     */
    public TreeGenerateUtil<I, O> objectTransformFunction(Function<? super I, ? extends O> objectTransform) {
        this.objectTransform = objectTransform;
        return this;
    }

    /**
     * 设置数据对象I中的主键字段
     *
     * @param keyMap 数据对象I中的主键函数
     * @return 返回this用于链式调用
     */
    public TreeGenerateUtil<I, O> keyMapFunction(Function<? super I, Object> keyMap) {
        this.keyMap = keyMap;
        return this;
    }

    /**
     * 设置 数据对象I中的标识层级关连的字段
     *
     * @param parentKeyMap 数据对象I中的标识层级关连的字段
     * @return 返回this用于链式调用
     */
    public TreeGenerateUtil<I, O> parentKeyMapFunction(Function<? super I, Object> parentKeyMap) {
        this.parentKeyMap = parentKeyMap;
        return this;
    }

    /**
     * 设置O对象中的子层级函数
     *
     * @param setChild 设置O对象中的子层级函数
     * @return 返回this用于链式调用
     */
    public TreeGenerateUtil<I, O> setChildFunction(BiConsumer<O, List<O>> setChild) {
        this.setChild = setChild;
        return this;
    }

    /**
     * 设置O对象排序比较函数
     *
     * @param sort 设置O对象中的子层级函数
     * @return 返回this用于链式调用
     */
    public TreeGenerateUtil<I, O> sortFunction(Comparator<? super I> sort) {
        this.sort = sort;
        return this;
    }

    /**
     * 生成树对象
     *
     * @param parentId 顶级的key
     * @return 树对象
     */
    public List<O> generate(Object parentId) {
        if (CollUtil.isEmpty(list)) {
            return Collections.emptyList();
        }
        Assert.notNull(objectTransform, "数据对象转化为tree对象的转化函数不能为null");
        Assert.notNull(keyMap, "数据对象获取主键映射函数不能为null");
        Assert.notNull(parentKeyMap, "数据对象获取标识层级关系属性映射函数不能为null");
        Assert.notNull(setChild, "设置tree对象中的子层级函数不能为null");
        return toTree(parentId);
    }

    private List<O> toTree(Object parentId) {
        Stream<I> stream = list.stream().filter(a -> Objects.equals(parentKeyMap.apply(a), parentId));
        if (Objects.nonNull(sort)) {
            stream = stream.sorted(sort);
        }
        List<I> parentList = stream.collect(Collectors.toList());
        if (CollUtil.isEmpty(parentList)) {
            return Collections.emptyList();
        }
        List<O> result = new ArrayList<>();
        parentList.forEach(p -> {
            Object key = keyMap.apply(p);
            if (handleKey.contains(key)) {
                return;
            }
            handleKey.add(key);
            List<O> childrenList = toTree(key);
            O node = objectTransform.apply(p);
            setChild.accept(node, childrenList);
            result.add(node);
        });
        return result;
    }
}
